import os
import vision.config as config
import numpy.random as npr

if __name__ == '__main__':

    net_postive_file = os.path.join(config.ANNO_STORE_DIR,config.ONET_POSTIVE_ANNO_FILENAME)
    net_part_file = os.path.join(config.ANNO_STORE_DIR,config.ONET_PART_ANNO_FILENAME)
    net_neg_file = os.path.join(config.ANNO_STORE_DIR,config.ONET_NEGATIVE_ANNO_FILENAME)
    net_ldsc_file = os.path.join(config.ANNO_STORE_DIR,config.ONET_CELEBA_LANDMARK_ANNO_FILENAME)
    net_ldsb_file = os.path.join(config.ANNO_STORE_DIR, config.ONET_BACKLIGHT_LANDMARK_ANNO_FILENAME)

    imglist_filename = config.ONET_TRAIN_IMGLIST_FILENAME

    anno_dir = config.ANNO_STORE_DIR
    imglist_file = os.path.join(anno_dir, imglist_filename)

    with open(net_postive_file, 'r') as f:
        pos = f.readlines()

    with open(net_part_file, 'r') as f:
        part = f.readlines()

    with open(net_neg_file, 'r') as f:
        neg = f.readlines()

    with open(net_ldsc_file, 'r') as f:
        ptsc = f.readlines()

    with open(net_ldsb_file, 'r') as f:
        ptsb = f.readlines()

    with open(imglist_file, "w") as f:

        pos_keep = npr.choice(len(pos),   size=min(len(pos),  1330000), replace=False)
        part_keep = npr.choice(len(part), size=min(len(part), 1330000), replace=False)
        neg_keep = npr.choice(len(neg),   size=min(len(neg),  4000000), replace=False)
        ptsc_keep = npr.choice(len(ptsc),   size=min(len(ptsc),  1330000), replace=False)
        ptsb_keep = npr.choice(len(ptsb), size=min(len(ptsb), 1330000), replace=False)

        print('pos num:{},choice:{}'.format(len(pos), len(pos_keep)))
        print('part num:{},choice:{}'.format(len(part), len(part_keep)))
        print('neg num:{},choice:{}'.format(len(neg), len(neg_keep)))
        print('celebA pts num:{},choice:{}'.format(len(ptsc), len(ptsc_keep)))
        print('backlight pts num:{},choice:{}'.format(len(ptsb), len(ptsb_keep)))

        for i in pos_keep:
            f.write(pos[i])
        for i in part_keep:
            f.write(part[i])
        for i in neg_keep:
            f.write(neg[i])
        for i in ptsc_keep:
            f.write(ptsc[i])
        for i in ptsb_keep:
            f.write(ptsb[i])

    print("ONet train annotation result file path:%s" % imglist_file)

